# 引入相关包
import os
import fileinput       #fileinput模块可以快速对一个或多个文件进行循环遍历
from neo4j import GraphDatabase
from config import NEO4J_CONFIG

driver = GraphDatabase.driver(**NEO4J_CONFIG)

def _load_data(path):
    """
    description: 将path目录下的csv文件以指定格式加载到内存
    :param path:  审核后的疾病对应症状的csv文件
    :return:      返回疾病字典，存储各个疾病以及与之对应的症状的字典
    """
    disease_dict=dict()
    disease_csv_list = os.listdir(path)
    # 初始化一个症状列表, 它里面是每种疾病对应的症状列表
    # 遍历疾病csv列表
    for disease_csv in disease_csv_list:
        disease = disease_csv.split(".")[0]  #将后缀.csv去掉, 获得疾病列表
        # 将疾病csv中的每个症状取出存入symptom列表中
        symptom_list = []
        with open(os.path.join(path, disease_csv), encoding='utf-8') as f:
            data_list = f.readlines()
            symptom_list.append([line.strip() for line in data_list if len(line) > 0])
        disease_dict[disease] = symptom_list

    return disease_dict     #{{'躁郁症'：'行为异常','非常烦躁','心神不定精神状态极差'},}

def write(path):
    """
    description: 将csv数据写入到neo4j, 并形成图谱
    :param path: 数据文件路径
    """
    # 使用_load_data从持久化文件中加载数据
    disease_symptom_dict = _load_data(path)
    print(disease_symptom_dict)
    # 开启一个neo4j的session
    with driver.session() as session:

        for key, value in disease_symptom_dict.items():
            cypher = "MERGE (a:Disease{name:%r}) RETURN a" % key
            session.run(cypher)
            for v in value:
                cypher = "MERGE (b:Symptom{name:%r}) RETURN b" % v
                session.run(cypher)
                cypher = "MATCH (a:Disease{name:%r}) MATCH (b:Symptom{name:%r}) WITH a,b MERGE(a)-[r:dis_to_sym]-(b)" % (key, v)
                session.run(cypher)
        cypher = "CREATE INDEX ON:Disease(name)"
        session.run(cypher)
        cypher = "CREATE INDEX ON:Symptom(name)"
        session.run(cypher)

# 输入参数path为csv数据所在路径
path = "C:/Users/fangjianwen/Desktop/structured/reviewed/"
#path = "C:/Users/fangjianwen/Desktop/test"
write(path)